﻿#pragma once
#include "yzrutil.h"
#include "lang/Thread.h"
#include <vector>
#include <queue>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>
#include <tuple>

namespace yzrilyzr_util{
	EBCLASS(FixedThreadPool){
	public:
	explicit FixedThreadPool(u_index threads) : stop(false){
		workers.reserve(threads);
		for(u_index i=0; i < threads; ++i){
			workers.emplace_back([this, i]{
				auto cur=yzrilyzr_lang::Thread::currentThread();
				cur->setName("FixedThreadPool Thread " + std::to_string(i));
				cur->setPriority(yzrilyzr_lang::Thread::Priority::HIGH);

				while(true){
					std::function<void()> task;

					{
						std::unique_lock<std::mutex> lock(queue_mutex);
						condition.wait(lock, [this]{
							return stop || !tasks.empty();
						});

						if(stop && tasks.empty()) return;

						task=std::move(tasks.front());
						tasks.pop();
					}

					try{
						task();
					} catch(...){
						std::lock_guard<std::mutex> lock(completion_mutex);
						--active_tasks;
						throw;
					}

					{
						std::lock_guard<std::mutex> lock(completion_mutex);
						--active_tasks;
						if(active_tasks == 0 && tasks.empty()){
							completion_condition.notify_all();
						} else{
							completion_condition.notify_one();
						}
					}
				}
			});
		}
	}

	template<class F, class... Args>
	auto commit(F && f, Args&&... args)
		-> std::future<std::invoke_result_t<F, Args...>>{
		using return_type=std::invoke_result_t<F, Args...>;

		auto task=std::make_shared<std::packaged_task<return_type()>>(
			[func=std::forward<F>(f),
			args_tuple=std::make_tuple(std::forward<Args>(args)...)]() mutable{
			return std::apply(func, std::move(args_tuple));
		}
		);

		std::future<return_type> res=task->get_future();

		{
			std::unique_lock<std::mutex> lock(queue_mutex);
			if(stop) throw std::runtime_error("enqueue on stopped ThreadPool");

			std::lock_guard<std::mutex> comp_lock(completion_mutex);
			++active_tasks;

			tasks.emplace([task](){ (*task)(); });
		}

		condition.notify_one();
		return res;
	}

	void waitAll(){
		std::unique_lock<std::mutex> lock(completion_mutex);
		completion_condition.wait(lock, [this](){
			return active_tasks == 0 && tasks.empty();
		});
	}

	~FixedThreadPool(){
		{
			std::unique_lock<std::mutex> lock(queue_mutex);
			stop=true;
		}
		condition.notify_all();
		for(std::thread & worker : workers){
			if(worker.joinable()){
				worker.join();
			}
		}
	}

	private:
	std::vector<std::thread> workers;
	std::queue<std::function<void()>> tasks;

	std::mutex queue_mutex;
	std::condition_variable condition;

	std::mutex completion_mutex;
	std::condition_variable completion_condition;
	std::atomic<u_index> active_tasks{0};  // 改为 atomic

	std::atomic<bool> stop;
	};
}